//
// Created by black on 2024/7/12.
//

#ifndef TDTRPROCESS_INTERRUPTIBLE_THREAD_H
#define TDTRPROCESS_INTERRUPTIBLE_THREAD_H

#include <thread>
#include <future>

class thread_interrupted : public std::exception{

};

class interrupt_flag {
    std::atomic<bool> flag;
    std::condition_variable* thread_cond = nullptr;
    std::condition_variable_any* thread_cond_any = nullptr;
    std::mutex set_clear_mutex;

    void interruption_point();

    template<class Lockable>
    struct custom_lock {
        interrupt_flag *self;
        Lockable &lk;
        custom_lock(interrupt_flag *self_, std::condition_variable_any &cond, Lockable &lk_) :
                self(self_), lk(lk_) {
            self->set_clear_mutex.lock();
            self->thread_cond_any = &cond;
        }

        void unlock() {
            lk.unlock();
            self->set_clear_mutex.unlock();
        }

        void lock() {
            std::lock(self->set_clear_mutex, lk);
        }

        ~custom_lock() {
            self->thread_cond_any = nullptr;
            self->set_clear_mutex.unlock();
        }
    };

public:
    void set();

    template<class Lockable>
    void wait(std::condition_variable_any& cv, Lockable& lk) {
        custom_lock cl(this, cv, lk);
        interruption_point();
        cv.wait(cl);
        interruption_point();
    }

    template<class Lockable, class Predicate>
    void wait(std::condition_variable_any& cv, Lockable& lk, Predicate&& pred) {
        custom_lock cl(this, cv, lk);
        interruption_point();
        cv.wait(cl,[&, this]{
            return !is_set() && !pred();
        });
        interruption_point();
    }

    template<class Lockable, class Predicate>
    void wait(std::condition_variable_any& cv, Lockable& lk, Predicate&& pred, int timeMillSeconds = 0) {
        custom_lock cl(this, cv, lk);
        interruption_point();
        if(timeMillSeconds > 0) {
            const auto timeout = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeMillSeconds);
            cv.wait_until(cl,timeout, [&, this]{
                return is_set() || pred();
            });
        }
        else {
            cv.wait(cl,[&, this]{
                return is_set() || pred();
            });
        }
        interruption_point();
    }

    bool is_set() const;
};

static thread_local interrupt_flag this_thread_interrupt_flag;

template<typename Lockable>
void interrupt_wait(std::condition_variable_any& cv, Lockable& lk) {
    this_thread_interrupt_flag.wait(cv, lk);
}

class interruptible_thread {
    std::thread internal_thread;
    interrupt_flag *flag = nullptr;

public:
    interruptible_thread() {

    }

    template<class FunctionType>
    explicit interruptible_thread(FunctionType f) {
        std::promise<interrupt_flag*> p;
        internal_thread = std::thread([f, &p] {
           p.set_value(&this_thread_interrupt_flag);
           try {
               f();
           }
           catch (const thread_interrupted& interrupted) {

           }
        });
        flag = p.get_future().get();
    }

    interruptible_thread(const interruptible_thread& other) = delete;
    interruptible_thread& operator=(const interruptible_thread& other) = delete;

    void interrupt() {
        if(flag) {
            flag->set();
        }
    }

    void join() {
        internal_thread.join();
    }

    bool joinable() const {
        return internal_thread.joinable();
    }

    void detach() {
        internal_thread.detach();
    }
};


#endif //TDTRPROCESS_INTERRUPTIBLE_THREAD_H
